package com.zym.init.async;

import cn.hutool.core.map.MapUtil;
import org.slf4j.MDC;
import org.springframework.core.task.TaskDecorator;
import org.springframework.lang.NonNull;

import java.util.Map;
import java.util.UUID;

import static com.zym.init.constant.TraceConstant.TRACE_ID;

/**
 * 追踪任务装饰器
 * <p>
 * 用于异步任务中为任务装饰追踪功能，保持任务执行过程中的追踪信息
 * <p>
 * 通过实现Spring的TaskDecorator接口来实现任务装饰器
 */
public class TraceTaskDecorator implements TaskDecorator {

    /**
     * 装饰器
     * 对传入的Runnable进行装饰，并返回装饰后的Runnable
     *
     * @param runnable 原始的{@code Runnable}
     * @return 装饰后的{@code Runnable}
     */
    @Override
    @NonNull
    public Runnable decorate(@NonNull Runnable runnable) {

        try {
            // 获取当前MDC的上下文信息
            Map<String, String> copyOfContextMap = MDC.getCopyOfContextMap();
            return () -> {
                try {
                    // 判断MDC上下文信息是否为空，若不为空则将上下文信息设置回MDC，否则设置新的追踪ID
                    if (MapUtil.isNotEmpty(copyOfContextMap)) {
                        MDC.setContextMap(copyOfContextMap);
                    } else {
                        MDC.put(TRACE_ID, UUID.randomUUID().toString().replace("-", ""));
                    }
                    // 执行传入的runnable
                    runnable.run();
                } finally {
                    // 清理MDC上下文信息
                    MDC.clear();
                }
            };
        } catch (IllegalStateException e) {
            // 如果在获取MDC上下文信息时出现异常，则直接返回原始的runnable
            return runnable;
        }
    }

}
